import sys, getopt, re, string, math
from array import array
from string import *

#global _project_name
#_project_name = "SORDO"

#global _file_geo_name
#_file_geo_name = _project_name.lower() + ".geo"

#global _file_geo
#_file_geo =  open(_file_geo_name, "w")

#global _lattice_file
#_lattice_file = _project_name.lower() + ".lattice"

#global _rotdefi_file
#_rotdefi_file = _project_name.lower() + ".rotdefi"

#global _rotdefi_i
#_rotdefi_i = int(1)

def CheckName(name):
    'checks BODY/REGION name for FLUKA com'
    if len(str(name)) > 8:
        print >>sys.stderr, "\033[31m WARNING: name", name, "is too long \033[0m"
        return 1
    if name[0] not in string.letters:
        print >>sys.stderr, "name:", name
        print >>sys.stderr, "the first character of body's name must be alphabetical"
        return 2
    return 0
        
class BODY:
    'class for a body'
    code = None
    name = None
    whats = []

    def __init__(self, name, w1, w2, w3, w4, w5, w6):
        self.name = name
        self.whats = w1, w2, w3, w4, w5, w6
#        CheckName(self.name)
    
    def Print(self, name, w1, w2, w3, w4, w5, w6, comment):
        CheckName(name)
        self.name = name
        if comment is None:
            print >> _file_geo, "%5s %8s %#9.4g %5.3f %5.3f %5.3f %5.3f %5.3f" %( self.code, name, w1, w2, w3, w4, w5, w6 )
        else:
            print >> _file_geo, "%5s %8s %#9.4g %#9.4g %#9.4g %#10.3f %#9.3f %#9.3f      ! %s" %( self.code, name, w1, w2, w3, w4, w5, w6, comment )

    def Print3(self, name, w1, w2, w3, comment):
        self.name = name
        print >> _file_geo, "%5s %8s %9.3f %9.3f %9.3f" % (self.code, self.name, w1, w2, w3)
#############################

class XYP(BODY):
    'XYP - plane perpendicular to the z-axis'
    z = None
    def __init__(self, name, z, comment=None):
        self.code = "XYP"
        self.name = name
        self.z = z
        self.Print(name, z, comment)

    def Print(self, name, z, comment):
        print >> _file_geo, "%5s %8s %9.3f" % (self.code, self.name, z)






_reg_counter = 0
global _materials
_materials = []

def ASSIGNMAT(fname=None):
    material_to_skip = "lattice"
    
    if fname is None:
        for i in range( len(_materials) ):
            if material_to_skip not in _materials[i]:
                print >> _file_geo, "ASSIGNMAT %16s" %_materials[i]
    else:
        file = open(fname, "w")
        file.write("! generated by pyfluka.py\n")
        for i in range( len(_materials) ):
            if material_to_skip not in _materials[i]:
                file.write("ASSIGNMAT %16s\n" %_materials[i])
        file.close()
    if len(_materials) != _reg_counter:
        print >>sys.stderr, "Warning: number of regions (%d) is not equal to one of materials (%d)" % (_reg_counter, len(_materials))

########################################

def LATTICE(container, cell, sdum=""):
    mode = "a"
    file = open(_lattice_file, mode)
    file.write("%-9.8s %-9.8s                     %-9.8s                       %-.8s\n" % ("LATTICE", container, cell, sdum))
    file.close()

        
def ROTDEFI(index, polar, azimutal, x, y, z):
    mode = "a"
#    if _rotdefi_i is 0: mode = "w"
    file = open(_rotdefi_file, mode)
#    if _rotdefi_i is 0: file.write("! generated by pyfiuka.py\n")
    file.write("%-9.8s  %9.1f %9.1f %9.1f %9.4f%9.4f %9.4f\n" % ("ROT-DEFI", index, polar, azimutal, x, y, z))
    file.close()

    
########################################

global _estimators
_estimators = []


class USRBIN:
    classname = 'USRBIN'
    type = None
    particle = None
    unit = None
    xmax = None
    ymax = None
    zmax = None
    name = None
    xmin = None
    ymin = None
    zmin = None
    nx   = None
    ny   = None
    nz   = None

    def __init__(self, type, particle, unit, xmax, ymax, zmax, name, xmin, ymin, zmin, nx, ny, nz):
        self.type = type
        self.particle = particle
        self.unit = unit
        self.xmax = xmax
        self.ymax = ymax
        self.zmax = zmax
        self.name = name
        CheckName(name)
        self.xmin = xmin
        self.ymin = ymin
        self.zmin = zmin
        self.nx = nx
        self.ny = ny
        self.nz = nz
        _estimators.append(self)

    def GetLine(self):
        return "%-9.8s %9.1f %9s %9.1f %9.3f %9.3f %9.3f %-.8s\n%-9.8s %9.3f %9.3f %9.3f %9.1f %9.1f %9.1f &\n*" % (self.classname, self.type, self.particle, self.unit, self.xmax, self.ymax, self.zmax, self.name, self.classname, self.xmin, self.ymin, self.zmin, self.nx, self.ny, self.nz)

    def Print(self):
        print >> _file_geo, self.GetLine()

class EVENTBIN(USRBIN):
    classname = 'EVENTBIN'

class USRYIELD:
    classname = 'USRYIELD'
    quantity = None
    particle = None
    unit = None
    upstream = None
    downstream = None
    norma = None
    sdum = None
    max1 = None
    min1 = None
    n1   = None
    max2 = None
    min2 = None
    type = None

    def __init__(self, quantity, particle, unit, upstream, downstream, norma, sdum, max1, min1, n1, max2, min2, type):
        self.quantity = quantity
        self.particle = particle
        self.unit = unit
        self.upstream = upstream
        self.downstream = downstream
        self.norma = norma
        self.sdum = sdum
        CheckName(sdum)
        self.max1 = max1
        self.min1 = min1
        self.n1 = n1
        self.max2 = max2
        self.min2 = min2
        self.type = type
        _estimators.append(self)

    def GetLine(self):
        return "%-9.8s %9.1f %9s %9.1f %9s %9s %9.1f %-.8s\n%-9.8s %9.3f %9.3f %9.1f %9.3f %9.3f %9.1f &\n*" % (self.classname, self.quantity, self.particle, self.unit, self.upstream, self.downstream, self.norma, self.sdum, self.classname, self.max1, self.min1, self.n1, self.max2, self.min2, self.type)
#        return "%-9.8s %9.1f %9s %9.1f %9s %9s %9.1f %-.8s \n*" % (self.classname, self.quantity, self.particle, self.unit, self.upstream, self.downstream, self.norma, self.sdum)

    def Print(self):
        print >> _file_geo, self.GetLine()
    
#############################

def LINE(fname, name, w1, w2, w3, w4, w5, w6, sdum):
    print >> _file_geo, "%s%5.1f%10s%5.1f%5.3f%5.3f%5.3f%s\n" % (name, w1, w2, w3, w4, w5, w6, sdum)

def PRINT_FILE(fname, where=sys.stdout):
    file = open(fname, "r")
    for line in file.readlines():
        line = line.rstrip()
        print >> where, line

def ESTIMATORS(fname=None):
    if fname is None:
        for i in range( len(_estimators) ):
            print >> _file_geo, _estimators[i].GetLine()
    else:
        file = open(fname, "w")
        file.write("! generated by pyfluka.py\n")
        for i in range( len(_estimators) ):
            file.write(_estimators[i].GetLine() + '\n')
        file.close()


#def RANDOMIZE(fname):
#    import random
#    file = open(fname, "w")
#    file.write("! generated by pyfluka.py\n")
#    file.write("%-9.8s %9.1f%9.0f.0\n" % ("RANDOMIZE", 1.0, random.random()*2e+6) )
#    file.close


class Project:
    name = None # name of the project (one word - used in the file names)
    title = None # title of the project
    geotitle = None # title of geometry
    inpfilename = None
    inpfile = None
    geofilename = None # name of the geometry file
    geofile = None # geometry file
    ivopt = None
    idbg = None
    region_counter = 0
    materials = []
    is_assignmat_printed = False # True if self.ASSIGNMAT() has been called. It must be done exactly once.
    regions = []

    def __init__(self, name, title, geotitle, inpfilename="", geofilename="", ivopt=0, idbg=0):
        self.name = name
        self.title = title
        self.CheckName(self.title, 80)
        self.geotitle = geotitle
        if len(inpfilename):
            self.inpfilename = inpfilename
        else:
            self.inpfilename = self.name.lower() + ".inp"
        self.inpfile = open(self.inpfilename, "w")
        if len(geofilename):
            self.geofilename = geofilename
        else:
            self.geofilename = self.name.lower() + ".geo"
        self.geofile = open(self.geofilename, "w")
        self.ivopt = ivopt
        self.idbg = idbg
#        self.region_counter = 0
        print >> self.inpfile, "TITLE\n%s" % self.title
        print >> self.inpfile, "*........+....1....+....2....+....3....+....4....+....5....+....6....+..SDUM.."
        print >> self.geofile, "%5d%5d               %s" % (self.ivopt, self.idbg, self.geotitle)


    def CheckName(self, name, maxl=8):
        'checks BODY/REGION name for FLUKA com'
        if len(str(name)) > maxl:
            self.error("CheckName: name %s is too long: %d symbols > %d" % (name, len(str(name)), maxl) )
            return True
        if name[0] not in string.letters:
            self.error("The first character of the body's name must be alphabetical (%s)" % name)
            return False
        return False

    def CheckRegion(self, name):
        "Checks whether the region name in the argument is in the list of regions"
        if name in self.regions:
            return True
        else:
            self.error("CheckRegion: there is no region called %s" % name)
            return False

    def end(self):
        'Prints END in the geometry file'
        print >> self.geofile, "  END"

    def warning(self, string):
        print >> sys.stderr, "\033[33mWARNING: %s\033[0m" % string

    def error(self, string):
        print >> sys.stderr, "\033[31mERROR: %s\033[0m" % string
        sys.exit(1)
        
    def close(self):
        self.inpfile.close()
        self.geofile.close()

    def ASSIGNMAT(self, fname=None):
        '''
        Print material assignment in the .inp file
        '''
        if self.is_assignmat_printed is True:
            self.error("Material assignment has already been printed in the input file. You should call Project.ASSIGNMAT() exactly once.")

        material_to_skip = "lattice"
    
        self.Comment("")
        for i in range( len(self.materials) ):
            if material_to_skip not in self.materials[i]:
                print >> self.inpfile, "ASSIGNMAT %16s" % self.materials[i]
        self.Comment("")

        if len(self.materials) != self.region_counter:
            self.warning("Number of regions (%d) is not equal to the number of materials (%d)" % (_reg_counter, len(self.materials)))

        self.is_assignmat_printed = True


    def Line(self, name, w1=None, w2=None, w3=None, w4=None, w5=None, w6=None, sdum=None, comment=None, format=None):
        '''
        Print a line in the input file according to the fixed FLUKA format.
        Alternatively, a user format can be specified if the last argument, 'format' is set.
        The fomat description used by Pytnon can be found here http://ru.wikipedia.org/wiki/Printf
        or, if you can read Russian, there is much more detailed description here: http://ru.wikipedia.org/wiki/Printf
        If 'comment' is set, it will be put as a comment line before the current one.
        '''
        if comment:
            self.Comment(comment)
        whats =  name, w1, w2, w3, w4, w5, w6, sdum
        if format is not None:
            print >> self.inpfile, format % whats
            return 0

        format = "%-9.9s"
        nonempty_whats = []
        nonempty_whats.append(name)
        for what in whats[1:7]:
            if what is None:
                format += " "*10
            else:
                nonempty_whats.append(what)
                try:
                    float(what)
                    format += "%#10.5G"
                except ValueError:
                    format += "%10s"
        if sdum is not None:
            format += " %s"
            nonempty_whats.append(sdum)
        print >> self.inpfile, format % tuple(nonempty_whats)

    def Open(self, unit, filename, state="NEW"):
        """
        Print line in the .inp file in order to open a unit with the given name
        """
        self.Line("OPEN", unit, None, None, None, None, None, state)
        self.RawLine(filename)

    def RawLine(self, line):
        '''
        Print an unformatted line in the input file
        '''
        print >> self.inpfile, line

    def Comment(self, comment):
        '''
        Print a comment line in the .inp file
        '''
        self.RawLine("* %s" % comment)

    def Print1(self, code, name, w1, comment=None):
        self.CheckName(name)
        print >> self.geofile, "%5s %8s %9.3f" % (code, name, w1)

    def Print3(self, code, name, w1, w2, w3, comment=None):
        self.CheckName(name)
        print >> self.geofile, "%5s %8s %9.3f %9.3f %9.3f" % (code, name, w1, w2, w3)

    def Print4(self, code, name, w1, w2, w3, w4):
        self.CheckName(name)
        print >> self.geofile, "%5s %8s %9.3f %9.3f %9.3f %9.3f" % (code, name, w1, w2, w3, w4)

    def Print6(self, code, name, w1, w2, w3, w4, w5, w6):
        self.CheckName(name)
        print >> self.geofile, "%5s %8s %#9.4f %9.4f %9.4f %9.4f %9.4f %9.4f" % ( code, name, w1, w2, w3, w4, w5, w6 )


    def Region(self, name, title, material="", comment=None, naz=5):
        'describes a region as a combination of bodies/zones'
        self.CheckName(name)
        self.region_counter = self.region_counter + 1
        if material is not "":
            self.CheckName(material)
            self.materials.append("%9s%11s" % (material, name))
        
        if comment is not None:
            print >> self.geofile,  '* region', self.region_counter, '-', comment
        words = title.split()
        len_of_line = 0
        for i in range(len(words)):
            len_of_line += len(words[i]) + 1
            if len_of_line > 70:
                words[i] = words[i] + "\n" + " "*12
                len_of_line = 0
#            if i and not i % 15: words[i] = words[i] + "\n" + " "*12
#            if i and not i % 13: words[i] = words[i] + "  "
        title = join(words)
        print >> self.geofile, "%9s %3d  %s" % (name, naz, title)
        self.regions.append(name)

    def SPH(self, name, x, y, z, R, comment=None):
        'SPH - sphere'
        self.Print4("SPH", name, x, y, z, R)

    def RPP(self, name, xmin, xmax, ymin, ymax, zmin=None, zmax=None, comment=None):
        '''
        RPP - rectangular parallelepiped
        It is possible to use both FLUKA and GEANT3 notations in constructor.
        FLUKA:  name, xmin,    xmax,    ymin,    ymax, zmin, zmax
        GEANT3: name, xcenter, ycenter, zcenter, array
        array is an array of 3 numbers of half-dimentions
        '''
        if zmax is not None: # FLUKA notation
            self.Print6("RPP", name, xmin, xmax, ymin, ymax, zmin, zmax)
            return BODY(name, xmin, xmax, ymin, ymax, zmin, zmax)
        else: # GEANT3 notation
            self.Print6("RPP", name, xmin-ymax[0], xmin+ymax[0], xmax-ymax[1], xmax+ymax[1], ymin-ymax[2], ymin+ymax[2])
            return BODY(name, xmin-ymax[0], xmin+ymax[0], xmax-ymax[1], xmax+ymax[1], ymin-ymax[2], ymin+ymax[2])


    def USRBIN(self, type, particle, unit, xmax, ymax, zmax, name, xmin, ymin, zmin, nx, ny, nz):
        self.CheckName(name)
        estname = "USRBIN"
        string =  "%-9.8s %9.1f %9s %9.1f %9.3f %9.3f %9.3f %-.8s\n%-9.8s %9.3f %9.3f %9.3f %9.1f %9.1f %9.1f &\n*" % (estname, type, particle, unit, xmax, ymax, zmax, name, estname, xmin, ymin, zmin, nx, ny, nz)
        print >> self.inpfile, string
        

    def USRBDX(self, quantity, particle, unit, upstream, downstream, norma, sdum, max1, min1, n1, max2, min2, type):
        self.CheckName(sdum)
        self.CheckRegion(upstream)
        self.CheckRegion(downstream)
        estname = "USRBDX"
        string = "%-9.8s %9.1f %9s %9.1f %9s %9s %9.1f %-.8s\n%-9.8s %10.3f %#9.3G%9.1f %9.3f %9.3f %9.1f &\n*" % (estname, quantity, particle, unit, upstream, downstream, norma, sdum, estname, max1, min1, n1, max2, min2, type)
        print >> self.inpfile, string

    def USRTRACK(self, type, particle, unit, detector, volume, nebins, name, emax, emin):
        self.CheckName(name)
        self.CheckRegion(detector)
        estname = "USRTRACK"
        string = "%-9.8s %9.1f %9s %9.1f %9s %9.1f %9.1f %-.8s\n%-8.8s %10.3f  %#9.4G %s &\n*" % (estname, type, particle, unit, detector, volume, nebins, name, estname, emax, emin, " "*39)
        print >> self.inpfile, string

            
    def XCC(self, name, y, z, R, comment=None):
        'XCC - cyllinder parallel to x-axis'
        self.Print3("XCC", name, y, z, R, comment)

    def XYP(self, name, z, comment=None):
        'XYP - plane perpendicular to the z-axis'
        self.Print1("XYP", name, z, comment)
 
    def XZP(self, name, y, comment=None):
        'XZP - plane perpendicular to the y-axis'
        self.Print1("XZP", name, y, comment)

    def YCC(self, name, x, z, R, comment=None):
        'YCC - cyllinder parallel to y-axis' 
        self.Print3("XCC", name, x, z, R, comment)

    def YZP(self, name, x, comment=None):
        'YZP - plane perpendicular to the x-axis'
        self.Print1("YZP", name, x, comment)

    def ZCC(self, name, x, y, R, comment=None):
        'ZCC - cyllinder parallel to z-axis'
        self.Print3("ZCC", name, x, y, R, comment)
